#include <iostream>
#include <eigen3/Eigen/Core>
#include <vector>
#include <functional>
#include <cmath>
typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> MatrixXd;
typedef Eigen::Matrix<double, Eigen::Dynamic, 1> VectorXd;
typedef std::function<double(double *)> Function;


enum Smoother
{
    WeightedJacobi = 1,
    GaussSeidel = 2
};

enum Restricter
{
    FullWeighting = 1,
    Injection = 2
};

enum Intper
{
    Linear = 1,
    Quadratic = 2
};

/**
 * @brief 多重网格求解器。
 */
class Multigrid_1D_Solver
{
protected:
    int _num_of_grid = 128;
    VectorXd _u_origin;
    VectorXd _f_origin;
    MatrixXd _BC;
    Smoother SmoothType = GaussSeidel;
    Restricter RestricType = FullWeighting;
    Intper IntpType = Linear;

public:
    VectorXd get_f_origin() { return _f_origin; };
    VectorXd get_u_origin() { return _u_origin; };
    double get_h() { return 1.0 / _num_of_grid; };
    void setSmoothType(int SmoothTypeCode) { SmoothType = (enum Smoother)SmoothTypeCode; };
    void setRestircType(int RestricTypeCode) { RestricType = (enum Restricter)RestricTypeCode; };
    void setIntpType(int IntpTypeCode) { IntpType = (enum Intper)IntpTypeCode; };
    void setNumOfGrid(int num) { _num_of_grid = num; };
    void Initialize_u();
    void Initialize_BC(Function _BCF);
    void InitializeRhs(Function _F);
    void Initialize(int num, Function _F, Function _BCF);
    VectorXd SmoothingforVC(VectorXd _u, VectorXd _f, double _h);
    VectorXd ResidualSolver(const VectorXd _u,const VectorXd _f,const double _h);
    VectorXd RestrictSolver(VectorXd _r);
    VectorXd IntpSolver(VectorXd _r);
    VectorXd Vcycle(VectorXd _u, VectorXd _f, double _h);
    VectorXd FMGcycle( VectorXd _f, double _h);
};

void Multigrid_1D_Solver::Initialize_BC(Function _BCF)
{
    double a[1]={0};
    _BC.setZero(2,1);
    _BC(0,0) = _BCF(a);
    a[0]=1;
    _BC(1,0)=_BCF(a);
};

void Multigrid_1D_Solver::Initialize(int num,Function _F, Function _BCF)
{
    setNumOfGrid(num);
    Initialize_BC(_BCF);
    InitializeRhs(_F);
    Initialize_u();
    std::cout << "Initialized!" << std::endl;
};

void Multigrid_1D_Solver::Initialize_u()
{   
    int N= _num_of_grid;
    _u_origin.setZero(1 + N);
    _u_origin(0) = _BC(0,0);
    _u_origin(N) = _BC(1,0);
};

void Multigrid_1D_Solver::InitializeRhs(Function _F)
{   
    int num= _num_of_grid;
    double a[1] = {0};
    _f_origin.resize(num + 1);
    for (int i = 1; i < num; i++)
    {
        a[0] = 1.0 * i / num;
        _f_origin(i) = _F(a);
    };
    _f_origin(0) = _BC(0,0);
    _f_origin(num) = _BC(1,0);
};

VectorXd Multigrid_1D_Solver::SmoothingforVC(VectorXd _u, VectorXd _f, double _h)
{
    VectorXd _u_temp;
    int N = _u.rows() - 1;
    _u_temp.setZero(N + 1);
    /// 由于 _u的首项与末项是ghost cell，所以在_u_temp中予以保留。
    _u_temp(0) = _u(0);
    _u_temp(N) = _u(N);
    double _t = 0;
    double omega = 2.0 / 3;
    switch (SmoothType)
    {
    case GaussSeidel:
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(i) = (_u_temp(i - 1) + _u(i + 1) - _h * _h * _f(i)) / 2.0;
        };
        break;
    case WeightedJacobi:
        for (int i = 1; i <= N - 1; i++)
        {
            _t = (_u(i - 1) + _u(i + 1) - _h * _h * _f(i)) / 2.0;
            _u_temp(i) = (1.0 - omega) * _u(i) + omega * _t;
        };
        break;
    default:
        break;
    };
    return _u_temp;
};

VectorXd Multigrid_1D_Solver::ResidualSolver(const VectorXd _u,const VectorXd _f,const double _h)
{
    VectorXd _u_temp;
    int N = _u.rows() - 1;
    _u_temp.setZero(N + 1);
    for (int i = 1; i <= N - 1; i++)
        _u_temp(i) = _f(i) - (_u(i - 1) - 2.0 * _u(i) + _u(i + 1)) / (_h * _h);
    return _u_temp;
};

VectorXd Multigrid_1D_Solver::RestrictSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (_r.rows() - 1) / 2;
    _u_temp.setZero(N + 1);
    switch (RestricType)
    {
    case FullWeighting:
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(i) = (_r(2 * i - 1) + 2.0 * _r(2 * i) + _r(2 * i + 1)) / 4.0;
        };
        break;
    case Injection:
        for (int i = 1; i <= N - 1; i++)
        {
            _u_temp(i) = _r(2 * i);
        };
        break;
    default:
        break;
    };
    _u_temp(0)=_r(0);
    _u_temp(N)=_r(2*N);
    return _u_temp;
};
VectorXd Multigrid_1D_Solver::IntpSolver(VectorXd _r)
{
    VectorXd _u_temp;
    int N = (_r.rows() - 1) * 2;
    _u_temp.setZero(N + 1);
    switch (IntpType)
    {
    case Linear:
        for (int i = 1; i <= N - 1; i = i + 2)
            _u_temp(i) = (_r((i - 1) / 2) + _r((i + 1) / 2)) / 2.0;
        for (int j = 0; j <= N; j = j + 2)
            _u_temp(j) = _r(j / 2);
        break;
    case Quadratic:
        _u_temp(1)= _r(0) / 4.0 + 5.0 * _r(1) / 8.0 + _r(2) / 8.0;
        for (int i = 3; i <= N - 3; i = i + 2)
        {
            _u_temp(i) = 3.0*_r((i - 1) / 2) / 16.0 + 5.0 * _r((i + 1) / 2) / 8.0 + 3.0*_r((i + 3) / 2) / 16.0;
        };
        _u_temp(N - 1) = _r(N / 2 - 2) / 8.0 + 5.0 * _r(N / 2 - 1) / 8.0 + _r(N / 2) / 4.0;
        for (int j = 0; j <= N; j = j + 2)
            _u_temp(j) = _r(j / 2);
        break;
    default:
        break;
    };
    _u_temp(0)=_r(0);
    _u_temp(N)=_r(N/2);
    return _u_temp;
};
VectorXd Multigrid_1D_Solver::Vcycle(VectorXd _u, VectorXd _f, double _h)
{
    VectorXd __u, __r, __eps;
    __u = SmoothingforVC(_u, _f, _h);
    __r = ResidualSolver(__u, _f, _h);
    __r = RestrictSolver(__r);
    __eps.setZero(__r.rows());
    if ((__eps.rows() - 1) == 2)
    {
  
        __eps = SmoothingforVC(__eps, __r, 2 * _h);
    }
    else
    {
        __eps = Vcycle(__eps, __r, 2 * _h);
    };
    __u = __u + IntpSolver(__eps);
    __u = SmoothingforVC(__u, _f, _h);
    return __u;
};

VectorXd Multigrid_1D_Solver::FMGcycle( VectorXd _f, double _h)
{
     VectorXd __u, __r;
     int N=_f.rows()-1;
    __u.setZero(N+1);
    __u(0)=_BC(0,0);
    __u(N)=_BC(1,0);
  
    if(N==2)
    {
        __u=SmoothingforVC(__u,_f,_h);
    }
    else
    {
        __r=RestrictSolver(_f);
        __u=FMGcycle(__r,2*_h);
        __u=IntpSolver(__u);
        __u=Vcycle(__u,_f,_h);
    };
    return __u;
};